import os
import dashscope
import clingo
from datetime import datetime

# 在代码最开头明确定义 ASP 代码（重要！）
asp_code = """
#const num_scenes = 7.
scene(1..num_scenes).
function(intro_char; intro_rival_char; add_conflict; add_bonding; add_obstacle_towards_major_goal; add_obstacle; add_breakthrough; add_twist; level_up_obstacle).

% 每个场景有一个功能
1 { scene_performs_func(S,F) : function(F) } 1 :- scene(S).

% 避免连续重复功能
:- scene_performs_func(S1,F), scene_performs_func(S2,F), S2 = S1 + 1.

% 冲突前至少两个角色
:- scene_performs_func(S2,add_conflict), 0 { scene_performs_func(S1,intro_char) : scene(S1), S1 < S2 } 1.

% 禁止在第4场景前添加转折点
:- scene_performs_func(S,add_twist), S < 4.
"""
# 在代码开头（建议在asp_code定义之后）明确定义翻译字典
translation_dict = {
    "intro_char": "Introduce a new character and describe their appearance and personality traits.",
    "intro_rival_char": "Introduce a character who is in opposition to the protagonist and explain the points of conflict between them.",
    "add_conflict": "Add a conflict event to the story to make the plot tense.",
    "add_bonding": "Describe a plot where two characters establish a certain connection or friendship.",
    "add_obstacle_towards_major_goal": "Set up an obstacle for the protagonist's main goal.",
    "add_obstacle": "Add an ordinary obstacle or difficulty.",
    "add_breakthrough": "The protagonist achieves a breakthrough when facing difficulties.",
    "add_twist": "Set up an unexpected twist to change the direction of the story.",
    "level_up_obstacle": "Upgrade the current obstacle to increase the difficulty of the story."
}

def print_log(message, level="INFO"):
    """统一的日志输出函数"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{level}] {message}")


def generate_outline(asp_content=None):
    """
    使用Clingo生成故事大纲（带自定义ASP代码支持）

    参数：
    asp_content: str - 可选的替代ASP代码

    返回：
    list - 生成的大纲功能序列
    """
    print_log("开始生成故事大纲...")

    # 确定使用的ASP代码
    used_asp_code = None
    try:
        # 优先级：参数 > 全局变量
        if asp_content:
            print_log("使用自定义ASP代码")
            used_asp_code = asp_content
        else:
            global asp_code  # 显式声明使用全局变量
            if 'asp_code' in globals():
                print_log("使用全局ASP代码")
                used_asp_code = asp_code
            else:
                print_log("错误：未找到任何ASP代码定义", "ERROR")
                return []

        # 验证代码有效性
        if not isinstance(used_asp_code, str) or len(used_asp_code.strip()) < 50:
            print_log(f"无效ASP代码（长度：{len(used_asp_code) if used_asp_code else 0}）", "ERROR")
            return []

    except Exception as e:
        print_log(f"初始化ASP代码时发生异常：{str(e)}", "ERROR")
        return []

    # Clingo处理流程
    control = clingo.Control()
    try:
        print_log("正在加载ASP代码到Clingo...")
        control.add("base", [], used_asp_code)

        print_log("基础规则初始化...")
        control.ground([("base", [])])

        models = []
        print_log("开始求解...")
        with control.solve(yield_=True) as handle:
            for model in handle:
                print_log(f"找到候选模型 #{len(models) + 1}")
                atoms = model.symbols(atoms=True)
                outline = []

                # 解析原子命题
                for atom in atoms:
                    if atom.name == "scene_performs_func":
                        try:
                            scene = atom.arguments[0].number
                            func = atom.arguments[1].name
                            outline.append((scene, func))
                        except Exception as e:
                            print_log(f"解析原子错误：{atom} - {str(e)}", "WARNING")

                # 排序并存储
                outline.sort(key=lambda x: x[0])
                models.append([func for _, func in outline])

                # 安全限制：最多取3个模型
                if len(models) >= 3:
                    print_log("达到最大模型限制（3个），停止求解")
                    break

        # 结果处理
        if models:
            print_log(f"共找到{len(models)}个有效模型")
            final_outline = models[0]  # 取第一个有效模型
            print_log(f"最终大纲顺序：{final_outline}")
            return final_outline
        else:
            print_log("未找到符合约束的模型", "WARNING")
            return []

    except RuntimeError as e:
        print_log(f"Clingo运行时错误：{str(e)}", "ERROR")
        return []
    except Exception as e:
        print_log(f"求解过程中发生未知异常：{str(e)}", "ERROR")
        return []

def translate_outline(outline):
    """将大纲中的功能翻译为写作指令"""
    print_log("开始翻译大纲为指令...")
    try:
        # 添加字典存在性检查
        if 'translation_dict' not in globals():
            print_log("错误：translation_dict未定义", "ERROR")
            return []

        translated = []
        for idx, func in enumerate(outline, 1):
            # 添加键存在性检查
            if func in translation_dict:
                print_log(f"场景{idx}: 转换功能 '{func}' => '{translation_dict[func]}'")
                translated.append(translation_dict[func])
            else:
                print_log(f"场景{idx}: 发现未知功能 '{func}'，已跳过", "WARNING")
        return translated
    except Exception as e:
        print_log(f"翻译大纲时发生异常：{str(e)}", "ERROR")
        return []

def generate_story(premise, instructions):
    """使用DashScope生成英语故事"""
    print_log("开始生成故事内容...")
    story = ""
    system_prompt = "You are a professional English story writer."

    if not instructions:
        print_log("没有可用的写作指令，跳过生成步骤", "WARNING")
        return ""

    for idx, instruction in enumerate(instructions, 1):
        print_log(f"正在处理第 {idx}/{len(instructions)} 个场景指令")
        print_log(f"当前指令: {instruction}")

        try:
            # 构建提示语
            if idx == 1:
                prompt = f"Premise: {premise}\nInstruction: {instruction}\nGenerate the first paragraph"
                print_log("生成首段提示语")
            else:
                prompt = f"Premise: {premise}\nInstruction: {instruction}\nCurrent story: {story[-500:]}\nContinue the next paragraph"  # 只保留最近500字符避免过长
                print_log("生成续写提示语（带上下文截断）")

            # 显示API调用参数（不显示敏感信息）
            print_log(f"调用API参数：model=qwen-max, prompt_length={len(prompt)}")

            # 检查API密钥
            if not os.getenv('DASHSCOPE_API_KEY'):
                print_log("未检测到DASHSCOPE_API_KEY环境变量！", "ERROR")
                return story + "\n[ERROR: MISSING API KEY]"

            # 调用API
            response = dashscope.Generation.call(
                api_key=os.getenv('DASHSCOPE_API_KEY'),
                model="qwen-max",
                messages=[
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': prompt}
                ],
                result_format='message'
            )

            # 处理响应
            print_log(f"收到API响应，状态码：{response.status_code}")
            if response.status_code == 200:
                if hasattr(response, 'output') and hasattr(response.output, 'choices'):
                    if len(response.output.choices) > 0:
                        content = response.output.choices[0].message.get('content', '')
                        print_log(f"成功获取内容，长度：{len(content)}字符")
                        story += content + "\n\n"
                        print_log(f"当前故事总长度：{len(story)}字符")
                    else:
                        print_log("API返回空的choices列表", "WARNING")
                        story += "[CONTENT MISSING]\n\n"
                else:
                    print_log("意外的API响应结构", "ERROR")
                    print_log(f"完整响应结构：{response}")
                    story += "[RESPONSE FORMAT ERROR]\n\n"
            else:
                print_log(f"API调用失败，状态码：{response.status_code}", "ERROR")
                print_log(f"错误详情：{response}")
                story += f"[API ERROR {response.status_code}]\n\n"

        except Exception as e:
            print_log(f"生成场景{idx}时发生异常：{str(e)}", "ERROR")
            story += f"[GENERATION ERROR: {str(e)}]\n\n"

        print_log("-" * 50)  # 分割线

    return story


if __name__ == "__main__":
    # 示例前提（英语）
    premise = """In a faraway kingdom, there lived a beautiful young girl named Ella. After her father's sudden death,
    she was forced to work as a servant for her wicked stepmother and two jealous stepsisters."""

    print_log("故事生成流程启动")
    print_log(f"初始前提：\n{premise}")

    # 生成大纲
    outline = generate_outline()
    print_log(f"原始大纲结构：{outline}")

    # 翻译大纲
    instructions = translate_outline(outline)
    print_log(f"最终写作指令：")
    for i, instr in enumerate(instructions, 1):
        print_log(f"指令{i}: {instr}")

    # 生成故事
    story = generate_story(premise, instructions)

    # 最终输出
    print_log("故事生成完成")
    print("\n" + "=" * 50 + " FINAL STORY " + "=" * 50)
    print(story)